package com.djhu.newscollector.crawel.net;

import com.djhu.newscollector.crawel.core.Url;
import okhttp3.*;

import javax.net.ssl.*;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.Charset;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.security.cert.CertificateFactory;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
 * Created by zw on 2018/2/11.
 */
public class OkHttpVisitor implements HttpVisitor{

    private static OkHttpClient okHttpClient;

    public static OkHttpClient getOkHttpClient(InputStream... certificates) {
        if (okHttpClient == null) {
            int cacheSize = 10 * 1024 * 1024;
            OkHttpClient.Builder builder = new OkHttpClient.Builder()
                    .connectTimeout(20, TimeUnit.SECONDS)
                    .writeTimeout(20, TimeUnit.SECONDS)
                    .readTimeout(20, TimeUnit.SECONDS);
            if (certificates != null){
                builder.sslSocketFactory(getSSLSocketFactory(certificates));
            }
            okHttpClient = builder.build();
        }
        return okHttpClient;
    }


    private static SSLSocketFactory getSSLSocketFactory(InputStream... certificates) {
        try {
            CertificateFactory certificateFactory = CertificateFactory.getInstance("X.509");
            KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
            keyStore.load(null);
            int index = 0;
            for (InputStream certificate : certificates) {
                String certificateAlias = Integer.toString(index++);
                keyStore.setCertificateEntry(certificateAlias, certificateFactory.generateCertificate(certificate));

                try {
                    if (certificate != null)
                        certificate.close();
                } catch (IOException e) {
                }
            }
            SSLContext sslContext = SSLContext.getInstance("TLS");
            TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            trustManagerFactory.init(keyStore);
            sslContext.init(null, trustManagerFactory.getTrustManagers(), new SecureRandom());
            return sslContext.getSocketFactory();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    private Set<Integer> successCodeSet;
    public OkHttpVisitor(){
        successCodeSet  =createSuccessCodeSet();
    }

    public Request.Builder createRequestBuilder(HttpRequest request) {
        Request.Builder builder = new Request.Builder()
                .header("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/60.0.3112.101 Safari/537.36")
                .url(request.getUrl().getUrl());
        return builder;
    }


    protected HashSet<Integer> createSuccessCodeSet(){
        HashSet<Integer> result = new HashSet<Integer>();
        result.add(200);
        result.add(301);
        result.add(302);
        result.add(404);
        return result;
    }

        @Override
    public Page visit(HttpRequest httpRequest) throws IOException {
        OkHttpClient okHttpClient = getOkHttpClient();
        Request  request = createRequestBuilder(httpRequest).build();

            Response response = okHttpClient.newCall(request).execute();

            String contentType = null;
            byte[] content = null;
            String charset = null;

            ResponseBody responseBody = response.body();
            int code = response.code();
            if(!successCodeSet.contains(code)){
//            throw new IOException(String.format("Server returned HTTP response code: %d for URL: %s (CrawlDatum: %s)", code,crawlDatum.url(), crawlDatum.key()));
                throw new IOException(String.format("Server returned HTTP response code: %d for %s", code,httpRequest.getUrl()));

            }
            if(responseBody != null){
                content = responseBody.bytes();
                MediaType mediaType = responseBody.contentType();
                if(mediaType!=null){
                    contentType = mediaType.toString();
                    Charset responseCharset = mediaType.charset();
                    if(responseCharset!=null){
                        charset = responseCharset.name();
                    }
                }
            }

            HttpResponse httpResponse = new HttpResponse();
            httpResponse.setCode(code);
            httpResponse.setRequest(httpRequest);
            httpResponse.setBytes(content);
            Page page = new Page();
            page.setResponse(httpResponse);
            page.charset(charset);
            page.obj(response);
            return page;
    }

    public static void main(String[] args) throws IOException {
        Url url = new Url();
        url.setUrl("https://www.s6tu.com/images/2018/01/04/1220_120.jpg");
        HttpRequest request =new HttpRequest();
        request.setUrl(url);

        OkHttpVisitor visitor = new OkHttpVisitor();
        Page page = visitor.visit(request);


    }
}